import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import prob
from NN import cdinn2


def plot_init():
    if prob.init_shape == 1: # rectangle
        init = matplotlib.patches.Rectangle((prob.init_min[0], prob.init_min[1]), \
            prob.init_max[0] - prob.init_min[0], prob.init_max[1] - prob.init_min[1], facecolor='green')
    if prob.init_shape == 2: # circle
        init = matplotlib.patches.Circle(((prob.init_min[0] + prob.init_max[0]) / 2.0, \
                (prob.init_min[1] + prob.init_max[1]) / 2.0), (prob.init_max[1] - prob.init_min[1]) / 2.0, facecolor='green')
    return init


def plot_unsafe(sub_unsafe, unsafe_shape):
    unsafe_min = sub_unsafe[0]
    unsafe_max = sub_unsafe[1]
    if unsafe_shape == 1:
        unsafe = matplotlib.patches.Rectangle((unsafe_min[0], unsafe_min[1]), \
            unsafe_max[0] - unsafe_min[0], unsafe_max[1] - unsafe_min[1], facecolor='red')
    elif unsafe_shape == 2:
        unsafe = matplotlib.patches.Circle(((unsafe_min[0] + unsafe_max[0]) / 2.0,
                (unsafe_min[1] + unsafe_max[1]) / 2.0), (unsafe_max[1] - unsafe_min[1]) / 2.0, facecolor='red')
    else:
        y = np.linspace(-np.sqrt(2), np.sqrt(2), 1000)
        x = - y ** 2
        unsafe = plt.fill(x, y, 'r')
    return unsafe


def gen_plot_data():
    plot_sample_x = torch.linspace(prob.domain_min[0], prob.domain_max[0], prob.PLOT_LEN_B[0])
    plot_sample_y = torch.linspace(prob.domain_min[1], prob.domain_max[1], prob.PLOT_LEN_B[1])
    grid_sample = [plot_sample_x, plot_sample_y]
    # grid_sample = torch.tensor(grid_sample)
    mesh = torch.meshgrid(grid_sample)
    flatten = [torch.flatten(mesh[i]) for i in range(len(mesh))] # flatten the list of meshes
    plot_data = torch.stack(flatten, 1) # stack the list of flattened meshes
    return plot_data


def plot_boundary(model): # barrier boundary: contour plotting
    barrier_plot_nn_input = gen_plot_data()
    # apply the nn model but do not require gradient
    with torch.no_grad():
        output = model(barrier_plot_nn_input)
        barrier_plot_nn_output = output.reshape(prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[0]) # y_size * x_size
    plot_Z = barrier_plot_nn_output.numpy()

    plot_sample_x = np.linspace(prob.domain_min[0], prob.domain_max[0], prob.PLOT_LEN_B[0])
    plot_sample_y = np.linspace(prob.domain_min[1], prob.domain_max[1], prob.PLOT_LEN_B[1])
    plot_X, plot_Y = np.meshgrid(plot_sample_x, plot_sample_y)
    #plt.contourf(plot_X, plot_Y, plot_Z, [0], color='k')
    barrier_contour = plt.contour(plot_X.T, plot_Y.T, plot_Z, [0], linewidths=1, colors='black')
    # plt.clabel(barrier_contour, fontsize=20, colors='b')
    return barrier_contour


def simulation(model):
    fig = plt.figure(
        figsize=(3, 3),
        dpi=200
    )
    ax = fig.add_subplot(1, 1, 1)

    # ==========================init and unsafe==============================
    init = plot_init()  # plot initial
    ax.add_patch(init)

    for i in range(len(prob.sub_unsafe)):
        unsafe = plot_unsafe(prob.sub_unsafe[i], prob.sub_unsafe_shape[i])  # plot unsafe
        ax.add_patch(unsafe)
    ax.set_aspect("equal")

    # ==================================0水平集==================================
    boundary = plot_boundary(model)

    # ==================================轨迹==================================
    x_data = np.linspace(prob.init_min[0], prob.init_max[0], 4)
    y_data = np.linspace(prob.init_min[1], prob.init_max[1], 4)

    s = np.array(np.meshgrid(x_data, y_data))
    b = s.reshape(-1, order='F')
    s = b.reshape(-1, 2)

    x = s[:, 0]
    y = s[:, 1]

    x0 = torch.tensor(s, dtype=torch.float32).to("cpu")
    print(x0.shape[0])

    x_pre = x0
    step_nums = 1000
    result = [x_pre]
    for i in range(step_nums):
        x_next = prob.vector_field(x_pre)
        result.append(x_next)
        x_pre = x_next

    result = torch.cat(result, 1).cpu().detach().numpy()

    for i in range(result.shape[0]):
        route = result[i, :].reshape((-1, 2))
        ax.scatter(route[:, 0], route[:, 1], alpha=1, s=1)

    plt.axis([prob.domain_min[0], prob.domain_max[0], prob.domain_min[1], prob.domain_max[1]])
    plt.show()


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('./model/cdinn_eg6_1_20_0.1_0.1_1_train.pt'), strict=True)
    simulation(model)
    pass